package top.scauwlt.registry.zk.util;

import lombok.extern.slf4j.Slf4j;
import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.imps.CuratorFrameworkState;
import org.apache.curator.framework.recipes.cache.PathChildrenCache;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheListener;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.zookeeper.CreateMode;
import top.scauwlt.config.RpcConfig;
import top.scauwlt.registry.RegistryClear;

import java.net.InetSocketAddress;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

/**
 * @Author wmm
 * @Date 2022/7/24 23:26
 */
@Slf4j
public class CuratorUtils implements RegistryClear {

    private static final int MAX_RETRIES = 3;
    /**
     * 重连时间间隔: BASE_SLEEP_TIME * Math.max(1, random.nextInt(1 << (retryCount + 1)))
     */
    private static final int BASE_SLEEP_TIME = 1000;

    public static final String ZK_REGISTER_ROOT_PATH = "/scau-rpc";

    private static final String DEFAULT_ZOOKEEPER_ADDRESS = "127.0.0.1:2181";

    private static final Map<String, List<String>> SERVICE_ADDRESS_MAP = new ConcurrentHashMap<>();

    private static final Set<String> REGISTERED_PATH_SET = ConcurrentHashMap.newKeySet();

    private static CuratorFramework zkClient;

    public CuratorUtils() {

    }

    /**
     * 创建临时节点
     *
     * @param client
     * @param path
     */
    public static void createEphemeralNode(CuratorFramework client, String path) {
        try {
            if (REGISTERED_PATH_SET.contains(path) || client.checkExists().forPath(path) != null) {
                log.info("The node is exist : [{}]", path);
            }
            client.create().creatingParentsIfNeeded().withMode(CreateMode.EPHEMERAL).forPath(path);
            log.info("The node is created successfully : [{}]", path);
            REGISTERED_PATH_SET.add(path);
        } catch (Exception e) {
            log.error("create ephemeral node for path [{}] fail", path);
        }
    }

    /**
     * 获取rpc服务相关的所有节点
     *
     * @param zkClient
     * @param rpcServiceName
     * @return
     */
    public static List<String> getChildrenNodes(CuratorFramework zkClient, String rpcServiceName) {
        if (SERVICE_ADDRESS_MAP.containsKey(rpcServiceName)) {
            return SERVICE_ADDRESS_MAP.get(rpcServiceName);
        }
        String path = ZK_REGISTER_ROOT_PATH + "/" + rpcServiceName;
        List<String> result = null;
        try {
            result = zkClient.getChildren().forPath(path);
            SERVICE_ADDRESS_MAP.put(rpcServiceName, result);
            registerWatcher(zkClient, rpcServiceName);
        } catch (Exception e) {
            log.error("get children nodes for path [{}] fail", path);
        }
        return result;
    }

    public static void clearRegistry(CuratorFramework zkClient, InetSocketAddress inetSocketAddress) {
        REGISTERED_PATH_SET.stream().parallel().forEach(path -> {
            try {
                if (path.endsWith(inetSocketAddress.toString())) {
                    zkClient.delete().forPath(path);
                }
            } catch (Exception e) {
                log.error("clear registry for path [{}] fail : {}", path, e.getMessage());
            }
        });
        log.info("All registered services on the server are cleared:[{}]", REGISTERED_PATH_SET.toString());
    }

    /**
     * 获取zk客户端
     *
     * @return
     */
    public static CuratorFramework getClient() {
        if (zkClient != null && zkClient.getState() == CuratorFrameworkState.STARTED) {
            return zkClient;
        }
        /*Properties properties = PropertiesFileUtil.readPropertiesFile(RpcConfigEnum.RPC_CONFIG_PATH.getPropertityValue());
        String zookeeperAddress = DEFAULT_ZOOKEEPER_ADDRESS;
        if (properties != null && properties.getProperty(RpcConfigEnum.ZK_ADDRESS.getPropertityValue()) != null) {
            zookeeperAddress = properties.getProperty(RpcConfigEnum.ZK_ADDRESS.getPropertityValue());
        }*/
        String zookeeperAddress = RpcConfig.getInstance().getRegistryAddress();
        RetryPolicy retryPolicy = new ExponentialBackoffRetry(BASE_SLEEP_TIME, MAX_RETRIES);
        zkClient = CuratorFrameworkFactory.builder()
                .connectString(zookeeperAddress)
                .retryPolicy(retryPolicy)
                .build();
        zkClient.start();
        try {
            if (!zkClient.blockUntilConnected(30, TimeUnit.SECONDS)) {
                throw new RuntimeException("zookeeper connect time out for address: " + zookeeperAddress);
            }
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
        return zkClient;
    }

    /**
     * 注册某个节点的子节点监听器，当子节点发生变化时，获取当前子节点列表，放入MAP中
     *
     * @param zkClient
     * @param rpcServiceName
     * @throws Exception
     */
    private static void registerWatcher(CuratorFramework zkClient, String rpcServiceName) throws Exception {
        String servicePath = ZK_REGISTER_ROOT_PATH + "/" + rpcServiceName;
        // @param cacheData为true时能缓存节点内容
        PathChildrenCache pathChildrenCache = new PathChildrenCache(zkClient, servicePath, true);
        PathChildrenCacheListener pathChildrenCacheListener = (curatorFramework, pathChildrenCacheEvent) -> {
            List<String> result = curatorFramework.getChildren().forPath(servicePath);
            SERVICE_ADDRESS_MAP.put(rpcServiceName, result);
        };
        pathChildrenCache.getListenable().addListener(pathChildrenCacheListener);
        pathChildrenCache.start();
    }

    @Override
    public void clearRegistry(InetSocketAddress address) {
        CuratorUtils.clearRegistry(CuratorUtils.getClient(), address);
    }
}

